/**
 * @file yolov4_tiny_cpu_node.cpp
 * @brief 
 * @author Linfu Wei (ghowoght@qq.com)
 * @version 1.0
 * @date 2021-05-05
 * 
 * @copyright Copyright (c) 2021  WHU-EIS
 * https://www.cnblogs.com/jack123/p/13615164.html
 * https://cloud.tencent.com/developer/article/1668133
 */

#include <iostream>
#include <vector>
#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <mutex>

#include<fstream>

using std::vector;
using cv::Mat;
std::mutex m;


cv::Mat frame;
bool new_frame_flag = false;

void img_callback(const sensor_msgs::Image::ConstPtr& image_msg)
{
    // std::unique_lock<std::mutex> mylock(m);
    cv::Mat frame_raw = cv::Mat(image_msg->height,image_msg->width,CV_MAKETYPE(CV_8U,3),const_cast<uchar*>(&image_msg->data[0]), image_msg->step);
    cv::cvtColor(frame_raw,frame_raw,cv::COLOR_BGR2RGB);
    // if(!new_frame_flag){
        frame_raw.copyTo(frame);
        new_frame_flag = true;
    // }    
}

int main(int argc, char** argv){
    ros::init(argc, argv, "yolov4_tiny_cpu_node");
    ros::NodeHandle nh;

    ros::Subscriber img_sub = nh.subscribe("/camera/color/image_raw", 1, img_callback);

    ros::Rate loop_rate(30);

    // std::string weights = "/home/ghowoght/workspace/model/0505/yolov3-tiny.weights";
    // std::string cfg = "/home/ghowoght/workspace/model/0505/yolov3-tiny.cfg";
    // std::string labelsPath = "/home/ghowoght/workspace/model/0505/coco.names";
    std::string weights = "/home/ghowoght/workspace/model/backup0504/yolov4-tiny_best.weights";
    std::string cfg = "/home/ghowoght/workspace/model/backup0504/yolov4-tiny.cfg";
    std::string labelsPath = "/home/ghowoght/workspace/model/backup0504/voc.names";
    cv::dnn::Net net = cv::dnn::readNetFromDarknet(cfg, weights);  
    net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
    net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
    std::vector<cv::String> outNames = net.getUnconnectedOutLayersNames();
    for (int i = 0; i < outNames.size(); i++) {
        printf("output layer name : %s\n", outNames[i].c_str());
    }


    std::vector<std::string> classNamesVec;
    std::ifstream classNamesFile(labelsPath);
    
    if (classNamesFile.is_open())
    {
        std::string className = "";
        while (std::getline(classNamesFile, className))
            classNamesVec.push_back(className);
    }


    while(ros::ok()){
        if(new_frame_flag)
        {
            auto start = cv::getTickCount();

            Mat inputBlob = cv::dnn::blobFromImage(frame, 1 / 255.F, cv::Size(416, 416), cv::Scalar(), true, false);
            
            net.setInput(inputBlob);
            // 检测
            std::vector<Mat> outs;
            net.forward(outs, outNames);
            // ROS_INFO("ok: %d", outs.size());

            vector<cv::Rect> boxes;
            vector<int> classIds;
            vector<float> confidences;
            for (size_t i = 0; i<outs.size(); ++i)
            {
                // detected objects and C is a number of classes + 4 where the first 4
                float* data = (float*)outs[i].data;
                for (int j = 0; j < outs[i].rows; ++j, data += outs[i].cols)
                {
                    Mat scores = outs[i].row(j).colRange(5, outs[i].cols);
                    cv::Point classIdPoint;
                    double confidence;
                    minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);
                    if (confidence > 0.2)
                    {
                        int centerX = (int)(data[0] * frame.cols);
                        int centerY = (int)(data[1] * frame.rows);
                        int width = (int)(data[2] * frame.cols);
                        int height = (int)(data[3] * frame.rows);
                        int left = centerX - width / 2;
                        int top = centerY - height / 2;

                        classIds.push_back(classIdPoint.x);
                        confidences.push_back((float)confidence);
                        boxes.push_back(cv::Rect(left, top, width, height));
                    }
                }
            }

            vector<int> indices;
            cv::dnn::NMSBoxes(boxes, confidences, 0.2, 0.2, indices);
            for (size_t i = 0; i < indices.size(); ++i)
            {
                int idx = indices[i];
                cv::Rect box = boxes[idx];
                cv::String className = classNamesVec[classIds[idx]];
                putText(frame, className.c_str(), box.tl(), cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(255, 0, 0), 2, 8);
                rectangle(frame, box, cv::Scalar(0, 0, 255), 2, 8, 0);
            }
            float fps = cv::getTickFrequency() / (cv::getTickCount() - start);
            float time = (cv::getTickCount() - start) / cv::getTickFrequency();
            std::ostringstream ss;
            ss << "FPS : "<< fps <<" detection time: " << time*1000 << " ms";
            putText(frame, ss.str(), cv::Point(20, 20), 0, 0.5, cv::Scalar(0, 0, 255));
            imshow("YOLOv4-Detections", frame);
            cv::waitKey(25);

            
            new_frame_flag = false;
        }

        ros::spinOnce();
        loop_rate.sleep();
    }

}

